import unittest
import pandas as pd
import os as files
from swarm_network import data

class TestDataOperations(unittest.TestCase):
    def test_config_loading(self):
        # Путь к вашему файлу config.yaml
        config_file_path = 'config.yaml'

        # Загрузка конфигурации
        config = data.read_config(config_file_path)

        # Проверка наличия необходимых параметров
        self.assertIn('objects_type', config)
        self.assertIn('num_points', config)
        self.assertIn('square_side', config)
        # Добавьте другие параметры

        # Проверка корректности считанных значений (например, что num_points - это положительное число)
        self.assertIsInstance(config['num_points'], int)
        self.assertGreater(config['num_points'], 0)
        # Добавьте другие проверки 

    def test_save_data(self):
        # Создаем тестовый DataFrame
        array = {'X': [10, 20, 30], 'Y': [40, 50, 60]}
        df = pd.DataFrame(array)
        directory = 'data'
        filename = 'test_file_save.csv'

        # Сохраняем данные в другой каталог и с другим именем файла
        data.save_data(df, directory, filename)

        # Проверяем, что файл существует
        self.assertTrue(files.path.exists(f'{directory}/{filename}'))

        # Проверяем, что файл не пустой
        self.assertGreater(files.path.getsize(f'{directory}/{filename}'), 0)

    def test_get_data(self):

        directory = 'data'
        filename = 'test_file_get.csv'
        # Создаем тестовый DataFrame
        array = {'A': [1, 2, 3], 'B': [4, 5, 6]}

        df = pd.DataFrame(array)

        # Сохраняем данные в CSV файл
        df.to_csv(f'{directory}/{filename}', index=False)

        # Получаем данные из файла
        result = data.get_data(directory, filename)

        # Проверяем, что результат является DataFrame
        self.assertIsInstance(result, pd.DataFrame)

        # Проверяем, что данные совпадают с исходными
        pd.testing.assert_frame_equal(result, df)

class TestGenerateData(unittest.TestCase):

    def test_generate_points_random(self):

        num_points = 100
        attitude = 1000
        square_side = 5000
        deviation_z = 50

        result = data.generate_points_random(num_points, square_side, attitude, deviation_z)

        self.assertIsInstance(result, pd.DataFrame)

        # Проверяем, что количество объектов соответствует заданному значению
        self.assertEqual(len(result), num_points)

        # Проверяем, что DataFrame содержит ожидаемые столбцы
        expected_columns = ['object_id', 'name', 'x', 'y', 'z']
        self.assertListEqual(list(result.columns), expected_columns)

        # Проверяем, что координаты X, Y, Z находятся в ожидаемом диапазоне
        self.assertTrue((result['x'] >= 0).all())
        self.assertTrue((result['x'] <= square_side).all())
        self.assertTrue((result['y'] >= 0).all())
        self.assertTrue((result['y'] <= square_side).all())
        self.assertTrue((result['z'] >= attitude-deviation_z).all())
        self.assertTrue((result['z'] <= attitude).all())
        
    def test_generate_points_grid(self):

        num_points = 36
        attitude = 500
        square_side = 12000
        deviation_z = 10
        deviation_x_y = 20

        result = data.generate_points_grid(num_points, square_side,
                                       attitude, deviation_x_y, deviation_z)

        self.assertIsInstance(result, pd.DataFrame)

        # Проверяем, что количество объектов соответствует заданному значению
        self.assertEqual(len(result), num_points)

        # Проверяем, что DataFrame содержит ожидаемые столбцы
        expected_columns = ['object_id', 'name', 'x', 'y', 'z']
        self.assertListEqual(list(result.columns), expected_columns)

        # Проверяем, что координаты X, Y, Z находятся в ожидаемом диапазоне
        self.assertTrue((result['x'] >= 0).all())
        self.assertTrue((result['x'] <= square_side + deviation_x_y).all())
        self.assertTrue((result['y'] >= 0).all())
        self.assertTrue((result['y'] <= square_side + deviation_x_y).all())
        self.assertTrue((result['z'] >= attitude-deviation_z).all())
        self.assertTrue((result['z'] <= attitude).all())

    def test_generate_points_rows(self):

        num_points = 100
        rows = 5
        step = 500
        attitude = 500
        deviation_z = 10
        deviation_x_y = 20

        result = data.generate_points_rows(num_points, rows, step, attitude,
                        deviation_x_y, deviation_z)

        self.assertIsInstance(result, pd.DataFrame)

        # Проверяем, что количество объектов соответствует заданному значению
        self.assertEqual(len(result), num_points)

        # Проверяем, что DataFrame содержит ожидаемые столбцы
        expected_columns = ['object_id', 'name', 'x', 'y', 'z']
        self.assertListEqual(list(result.columns), expected_columns)
        self.assertTrue((result['z'] >= attitude-deviation_z).all())
        self.assertTrue((result['z'] <= attitude).all())

if __name__ == '__main__':
    unittest.main()
